import subprocess
import socket
import time
import logging
import sys

# 域名列表
domain_names = ["yourdomain"]  # 添加您要监测的域名
port = 60002  # change this

# IP地址列表
ip_addresses_to_allow = ["192.168.1.100", "192.168.1.101"]  # 添加要允许的IP地址

# 配置日志记录器
logger = logging.getLogger('updateUFW')
logger.setLevel(logging.INFO)

# 创建一个处理程序来输出到控制台
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)

# 创建一个处理程序来发送到系统日志
systemd_handler = logging.FileHandler('/var/log/updateUFW.log')  # 修改日志文件路径
systemd_handler.setLevel(logging.INFO)

# 配置日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)
systemd_handler.setFormatter(formatter)

# 添加处理程序到日志记录器
logger.addHandler(console_handler)
logger.addHandler(systemd_handler)

# 定义函数来解析域名并获取IP地址
def get_ip_from_domain(domain_name):
    try:
        ip_address = socket.gethostbyname(domain_name)
        return ip_address
    except Exception as e:
        logger.error(f"解析域名时发生错误: {str(e)}")
        return None

# 定义函数来获取所有UFW规则中允许访问指定端口的IP地址
def get_all_allowed_ips(port):
    try:
        ufw_rules = subprocess.check_output(f"sudo ufw show added | grep 'to any port {port}'", shell=True).decode("utf-8")
        lines = ufw_rules.splitlines()
        ufw_ips = []
        for line in lines:
            parts = line.split()
            if "ALLOW" in parts:
                ufw_ips.append(parts[parts.index("ALLOW") - 1])
        return ufw_ips
    except Exception as e:
        logger.error(f"获取UFW规则时发生错误: {str(e)}")
        return []

# 定义函数来更新UFW规则
def update_ufw_rules(domain_names, port, ip_addresses_to_allow):
    ufw_ips = get_all_allowed_ips(port)

    for domain_name in domain_names:
        ip_new = get_ip_from_domain(domain_name)
        if ip_new is not None:
            ip_addresses_to_allow.append(ip_new)

    # 删除不再允许的IP地址规则
    for ip in ufw_ips:
        if ip not in ip_addresses_to_allow:
            try:
                subprocess.run(f"sudo ufw delete allow from {ip} to any port {port}", shell=True, check=False)
                logger.info(f"已删除UFW规则允许IP地址 {ip} 访问端口 {port}")
            except Exception as e:
                logger.error(f"删除UFW规则时发生错误 for IP地址 {ip}: {str(e)}")

    # 添加新的IP地址规则
    for ip in ip_addresses_to_allow:
        if ip not in ufw_ips:
            try:
                subprocess.run(f"sudo ufw allow from {ip} to any port {port}", shell=True, check=True)
                logger.info(f"已添加UFW规则允许IP地址 {ip} 访问端口 {port}")
            except Exception as e:
                logger.error(f"添加UFW规则时发生错误 for IP地址 {ip}: {str(e)}")

# 定义定时任务，每隔一段时间执行一次
while True:
    update_ufw_rules(domain_names, port, ip_addresses_to_allow)
    # 等待10分钟再次执行
    time.sleep(600)
